Skip to main content

Prefix Sum

https://leetcode.com/problems/maximum-size-subarray-sum-equals-k/

  • [[Cheap Trick#^d4096b]]: Calculate subarray sum using prefix
def maxSubArrayLen(self, nums: List[int], k: int) -> int:
    prefix = [0] * len(nums)
    prefixIndex = collections.defaultdict(list)
    current = 0
    for i in range(len(nums)):
        prefix[i] = current
        prefixIndex[current].append(i)
        current += nums[i]

    res = 0
    for i in range(len(nums)):
        curSum = prefix[i] + nums[i]
        curDiff = curSum - k

        if curDiff in prefixIndex:
            for j in prefixIndex[curDiff]:
                if j <= i:
                    res = max(res, i -j + 1)
                    break

    return res

https://leetcode.com/problems/continuous-subarray-sum

def checkSubarraySum(self, nums: List[int], k: int) -> bool:
    prefixes = collections.defaultdict(int)
    prefixes[0] = -1

    prefix = 0
    for i, num in enumerate(nums):
        prefix += num
        rem = prefix % k
        if rem in prefixes:
            startIndex = prefixes[rem]
            if i - startIndex >= 2:
                return True
        else:
            prefixes[rem] = i
    
    return False

https://leetcode.com/problems/minimum-number-of-operations-to-move-all-balls-to-each-box/

  • At each index, number of moves increase by the number of balls 
def minOperations(self, boxes: str) -> List[int]:
    prefixBalls = 0
    prefixMove = 0
    res = [0] * len(boxes)

    for i in range(len(boxes)):
        res[i] += prefixMove
        if boxes[i] == "1":
            prefixBalls += 1
        prefixMove += prefixBalls
    
    postfixBalls = 0
    postFixMove = 0

    for j in range(len(boxes) -1, -1, -1):
        res[j] += postFixMove
        if boxes[j] == "1":
            postfixBalls += 1
        postFixMove += postfixBalls
    
    return res

https://leetcode.com/problems/sum-of-distances

  • Important
  def distance(self, nums: List[int]) -> List[int]:
    numIndexes = collections.defaultdict(list)

    for i, num in enumerate(nums):
        if num not in numIndexes:
            numIndexes[num] = [(i, i)]
        else:
            numIndexes[num].append((i, numIndexes[num][-1][1] + i))

    res = [0] * len(nums)
    for num, indexes in numIndexes.items():
        for i in range(len(indexes)):
            index, prefix = indexes[i]
            lenLeft, sumLeft = i, prefix - index
            lenRight, sumRight = len(indexes) - 1 - i, indexes[-1][1] - prefix

            current = index * lenLeft - sumLeft + sumRight - index * lenRight
            res[index] = current

    return res